import copy

# De Brujin 图

def get_reverse_complement(seq: str):
    pair = { 'A': 'T', 'T': 'A', 'G': 'C', 'C': 'G' }
    seq = seq[::-1]
    result = ''
    for i in seq:
        result.append(pair[i])
    return result

class Vertex():
    '''
    图中节点
    '''
    def __init__(self, s: str):
        self.s = s
        self.weight = 0
        # 出边
        self.outs = set()
        self.visited = False
        self.depth = 0
        self.max_depth_edge = None

    def reset(self):
        self.visited = False
        self.depth = 0
        self.max_depth_edge = None
    
    def add_edge(self, s):
        self.outs.add(s)

    def delete_edges(self, s):
        self.outs = self.outs - s


class Graph():
    '''
    De Brujin 图
    '''

    def __init__(self, dataset, config):
        '''
        dataset: reads 数据集
        config:
            length:   int    # k-mer 片段长度
            rev_comp?: bool   # （可选）是否加入反向互补序列，默认不加入
        '''
        self.k = config['length']
        self.vertices = {}
        # count 节点个数
        self.count = 0
        self.s2id = {}
        for i in dataset:
            self.read(i)
            if 'rev_comp' in config and config['rev_comp']:
                self.read(get_reverse_complement(i))

    def read(self, s: str):
        '''
        读取一个 read 序列
        '''
        last = self.add_vertex(s[0:self.k])
        for i in range(len(s) - self.k - 1):
            cur = self.add_vertex(s[i + 1: i + self.k + 1])
            self.add_edge(last, cur)
            last = cur

    def add_vertex(self, s: str):
        '''
        添加一个节点到图中（存在则 weight + 1）
        '''
        if s not in self.s2id:
            self.s2id[s] = self.count
            self.vertices[self.count] = Vertex(s)
            self.count += 1
        id = self.s2id[s]
        self.vertices[id].weight += 1
        return id

    def add_edge(self, u: int, v: int):
        '''
        添加一条边到图中
        u, v 为 Vertex 的 id
        '''
        self.vertices[u].add_edge(v)

    def get_weight(self, id):
        return self.vertices[id].weight
    
    def sort_edges(self, id: int):
        edges = list(self.vertices[id].outs)
        if len(edges) == 0:
            return []
        # edges.sort(key=lambda x: self.vertices[x].weight, reverse=True)
        edges.sort(key=lambda x: self.vertices[x].weight, reverse=True)
        return edges

    def reset(self):
        for i in self.vertices.values():
            i.reset()

    def get_longest_path(self):
        self.reset()
        max_depth = 0
        max_vertex = None
        for id in self.vertices.keys():
            depth = self.get_depth(id)
            if depth > max_depth:
                max_depth = depth
                max_vertex = id
        if None == max_vertex:
            return None
        path = []
        edge = max_vertex
        while edge:
            path.append(edge)
            edge = self.vertices[edge].max_depth_edge
        return path

    def get_depth(self, id):
        if not self.vertices[id].visited:
            self.vertices[id].visited = True
            edges = self.sort_edges(id)
            max_depth, max_edge = 0, None
            for edge in edges:
                depth = self.get_depth(edge)
                if depth > max_depth:
                    max_depth, max_edge = depth, edge
            self.vertices[id].depth, self.vertices[id].max_depth_edge = max_depth + 1, max_edge
        return self.vertices[id].depth

    def remove_longest_path(self, path):
        for i in path:
            del self.vertices[i]
        s = set(path)
        for i in self.vertices.values():
            i.delete_edges(s)

    def check(self):
        for i in self.vertices.values():
            a = set()
            for j in i.outs:
                assert(j.v not in a)
                a.add(j.v)

    def path2str(self, path):
        if len(path) < 1:
            return None
        s = self.vertices[path[0]].s
        for i in range(1, len(path)):
            s += self.vertices[path[i]].s[-1]
        return s

